# 前景主体分割算法 - 算法推理器
from networks.r2net import build_model
from cv2 import resize, IMREAD_GRAYSCALE, bitwise_and,imencode,imdecode,split,merge
from torch import nn, load, no_grad, sigmoid, device
from time import time
#import numpy as np
from numpy import fromfile,squeeze,uint8,asarray
def save_salmap(path, size, pred):
    pred = squeeze(sigmoid(pred).cpu().data.numpy())
    multi_fuse = 255 * pred
    multi_fuse = resize(multi_fuse, (int(size[0]), int(size[1])))
    #imwrite(path, multi_fuse)
    #imencode('.jpg',path)[1].tofile(multi_fuse)
    imencode('.png', multi_fuse)[1].tofile(path)


# 智能推理器
class Solver(object):
    def __init__(self, config):
        self.test_loader = None
        self.config = config
        self.build_network()

    # 构建网络模型
    def build_network(self):
        self.net = build_model(self.config.arch)
        if self.config.cuda:
            self.net = self.net.cuda()
        self.print_network(self.net, "R2Net Structure")
        self.load_weights()

    # 输出模型信息
    def print_network(self, model, name):
        num_params = 0
        for p in model.parameters():
            num_params += p.numel()
        print(name)
        print(model)
        print("The number of parameters: {}".format(num_params))

    # 加载模型参数
    def load_weights(self):
        print("Loading pre-trained model from %s..." % self.config.model)
        if len(self.config.test_on_gpus) > 0:
            self.net = nn.DataParallel(self.net, device_ids = self.config.test_on_gpus)

        if self.config.cuda:
            self.net.load_state_dict(load(self.config.model))
        else:
            self.net.load_state_dict(load(self.config.model, map_location = device('cpu')))
        self.net.eval()

    # 设置图片读取的 dataloader
    def setup_dataloader(self, test_loader):
        self.test_loader = test_loader


    # 生成 mask 结果图片
    def generate_mask(self, mask_path):
        img_num = len(self.test_loader)
        time_start = time()

        for i, data_batch in enumerate(self.test_loader):
            print("Testing %d/%d" % (i + 1, img_num))
            images, name, im_size = data_batch['image'], data_batch['name'][0], data_batch['size']

            with no_grad():
                global_pre, pre0, pre1, pre2, pre3, pre4 = self.net(images)
                save_salmap(mask_path, im_size, pre4)

        time_end = time()
        print("Speed: %f FPS" % (img_num / (time_end - time_start)))
        print("Test Finished !")


    # 生成 result 主体分割图片
    def generate_result(self, image_path, mask_path, result_path):
        #img_copy = imread(image_path)
        img_copy =imdecode(fromfile(image_path, dtype=uint8), -1)
        #mask = imread(mask_path, IMREAD_GRAYSCALE)
        mask = imdecode(fromfile(mask_path, dtype=uint8), IMREAD_GRAYSCALE)

        for height in range(mask.shape[0]):
            for length in range(mask.shape[1]):
                if mask[height][length] < 127:    mask[height][length] = 0
                if mask[height][length] > 127:    mask[height][length] = 255
        image = bitwise_and(img_copy, img_copy, mask = mask)

        for height in range(image.shape[0]):
            for length in range(image.shape[1]):
                if image[height][length][0] == 0 and image[height][length][1] == 0 and image[height][length][2] == 0:
                    image[height][length][0] = 255
                    image[height][length][1] = 255
                    image[height][length][2] = 255
        #imwrite(result_path, image)
        imencode('.png', image)[1].tofile(result_path)